package com.whirl.ai_ml;

import cn.hutool.core.util.NumberUtil;
import ml.dmlc.xgboost4j.java.Booster;
import ml.dmlc.xgboost4j.java.DMatrix;
import ml.dmlc.xgboost4j.java.XGBoost;
import ml.dmlc.xgboost4j.java.XGBoostError;

import java.math.BigDecimal;
import java.util.HashMap;
import java.util.Map;

import static com.whirl.util.Constants.MLData_PATH;

/**
 * @author xw
 * @date 2022/8/17 9:52
 */
public class XGBoostLoadRate {
    public static void main(String[] args) throws XGBoostError {

        DMatrix trainMat = new DMatrix(MLData_PATH + "load_rate/ai_ai.txt.train");
        DMatrix testMat = new DMatrix(MLData_PATH + "load_rate/ai_ai.txt.test");
//        DMatrix trainMat = new DMatrix(MLData_PATH + "load_rate/ai_d9e8.txt.train");  // 和总的差不多
//        DMatrix testMat = new DMatrix(MLData_PATH + "load_rate/ai_d9e8.txt.test");
//        DMatrix trainMat = new DMatrix(MLData_PATH + "load_rate/ai_c66.txt.train");  // 10%的误差，较大，因为数据少，只有78条
//        DMatrix testMat = new DMatrix(MLData_PATH + "load_rate/ai_c66.txt.test");

        System.out.println("训练集样本个数：" + trainMat.getLabel().length);

        // 获取预测标签label的值
        float[] originLabel = testMat.getLabel();
        System.out.println("测试集样本个数：" + originLabel.length);
        System.out.print("实际值： ");
        for (float i : originLabel) {
            System.out.print(i + "  ");
        }
        System.out.println();

        // 训练模型的参数
        HashMap<String, Object> params = new HashMap<String, Object>() {
            {
                put("silent", 1);
                put("objective", "reg:squarederror");
                put("booster", "gblinear");
                put("max_depth", 2);   // 最多分裂几次
                put("eta", 0.05);       // 学习率，越小时间越长
                put("gamma", 0.0005);    // 越大越保守（分裂次数少）
                put("subsample", 1); // 随机采样比例，抽取数据进行训练，控制过拟合
                put("colsample_bytree", 1);  // 抽取特征进行训练，控制变量间相关系数，来控制过拟合
                // 正则化参数，控制模型复杂度，来控制过拟合
                put("alpha", 0);    // L1正则化系数，越大越保守（分裂次数少）
                put("lambda", 1);  // L2正则化系数，越大越保守（分裂次数少）
            }
        };

        // 监督学习
        HashMap<String, DMatrix> watches = new HashMap<String, DMatrix>();
        watches.put("train", trainMat);
        watches.put("test", testMat);

        // 树的个数，开始训练
        int round = 400;
        Booster booster = XGBoost.train(trainMat, params, round, watches, null, null);

        // 保存模型
        booster.saveModel("model/xgb_model.model");
        System.out.println("模型保存成功");

        // 预测结果
        System.out.print("预测值： ");
        float[][] predicts = booster.predict(testMat);
        for (int i = 0; i < predicts.length; i++) {
            for (int j = 0; j < predicts[0].length; j++) {
                BigDecimal retain2 = NumberUtil.round(predicts[i][j], 2);
                predicts[i][j] = retain2.floatValue();
                System.out.print(predicts[i][0] + "  ");  // 写[i][j]结果也一样的
            }
        }
        System.out.print("\n误差：   ");

        // 预测值和实际值的差的绝对值
        float abs_error = 0;
        for (int i = 0; i < predicts.length; i++) {
            float err = Math.abs(predicts[i][0] - originLabel[i]);
            err = NumberUtil.round(err, 2).floatValue();
            System.out.print(err + "  ");
            abs_error += err;
        }
//        System.out.println("\n总误差： " + abs_error);
        System.out.println("\n平均误差： " + abs_error / predicts.length * 100 + "%");

    }
}
